import numpy as np
import torch

import numpy as np
import torch


class ResidualCoder(object):
    def __init__(self, code_size=8):
        super().__init__()
        self.code_size = code_size

    def encode_torch(self, boxes):
        """
        Args:
            boxes: (N, 7) [x, y, z, azimuth, elevation, grasp_angle, width]
        Returns:
        """
        boxes[:, 3:6] = torch.clamp_min(boxes[:, 3:6], min=1e-5)

        x_g, y_g, z_g, azimuth_g, elevation_g, grasp_angle_g, width_g, score = torch.split(boxes, 1, dim=-1)
        #diagonal = torch.sqrt(dxa ** 2 + dya ** 2)
        xt = x_g / 0.005
        yt = y_g / 0.005
        zt = z_g / 0.005
        azimuth_t = azimuth_g / 10 
        elevation_t = elevation_g / 10
        grasp_angle_t = grasp_angle_g / 10
        width_t = width_g / 0.005

        return torch.cat([xt, yt, zt, azimuth_t, elevation_t, grasp_angle_t, width_t, score], dim=-1)

    def decode_torch(self, box_encodings):
        """
        Args:
            box_encodings: (B, N, 7 + C) or (N, 7 + C) [x, y, z, dx, dy, dz, heading or *[cos, sin], ...]
        Returns:
        """
        xt, yt, zt, azimuth_t, elevation_t, grasp_angle_t, width_t, score = torch.split(box_encodings, 1, dim=-1)

        xg = xt * 0.005
        yg = yt * 0.005
        zg = zt * 0.005

        azimuth_g = azimuth_t * 10
        elevation_g = elevation_t * 10
        grasp_angle_g = grasp_angle_t * 10
        width_g = width_t * 0.005

        return torch.cat([xg, yg, zg, azimuth_g, elevation_g, grasp_angle_g, width_g, score], dim=-1)
